home *** CD-ROM | disk | FTP | other *** search
/ Chip 2006 June / CHIP 2006-06.2.iso / program / freeware / Democracy-0.8.2.exe / xulrunner / python / BitTorrent / StorageWrapper.py < prev    next >
Encoding:
Python Source  |  2006-04-10  |  16.0 KB  |  409 lines

  1. # The contents of this file are subject to the BitTorrent Open Source License
  2. # Version 1.0 (the License).  You may not copy or use this file, in either
  3. # source code or executable form, except in compliance with the License.  You
  4. # may obtain a copy of the License at http://www.bittorrent.com/license/.
  5. #
  6. # Software distributed under the License is distributed on an AS IS basis,
  7. # WITHOUT WARRANTY OF ANY KIND, either express or implied.  See the License
  8. # for the specific language governing rights and limitations under the
  9. # License.
  10.  
  11. # Written by Bram Cohen
  12.  
  13. from __future__ import division
  14.  
  15. from sha import sha
  16. from array import array
  17. from binascii import b2a_hex
  18.  
  19. from BitTorrent.bitfield import Bitfield
  20. from BitTorrent import BTFailure, INFO, WARNING, ERROR, CRITICAL
  21.  
  22. def toint(s):
  23.     return int(b2a_hex(s), 16)
  24.  
  25. def tobinary(i):
  26.     return (chr(i >> 24) + chr((i >> 16) & 0xFF) +
  27.         chr((i >> 8) & 0xFF) + chr(i & 0xFF))
  28.  
  29. NO_PLACE = -1
  30.  
  31. ALLOCATED = -1
  32. UNALLOCATED = -2
  33. FASTRESUME_PARTIAL = -3
  34.  
  35. class StorageWrapper(object):
  36.  
  37.     def __init__(self, storage, config, hashes, piece_size, finished,
  38.             statusfunc, flag, data_flunked, infohash, errorfunc, resumefile):
  39.         self.numpieces = len(hashes)
  40.         self.storage = storage
  41.         self.config = config
  42.         check_hashes = config['check_hashes']
  43.         self.hashes = hashes
  44.         self.piece_size = piece_size
  45.         self.data_flunked = data_flunked
  46.         self.errorfunc = errorfunc
  47.         self.total_length = storage.get_total_length()
  48.         self.amount_left = self.total_length
  49.         self.partial_mark = "BitTorrent - this part has not been "+\
  50.                             "downloaded yet."+infohash+\
  51.                             tobinary(config['download_slice_size'])
  52.         if self.total_length <= piece_size * (self.numpieces - 1):
  53.             raise BTFailure, 'bad data in responsefile - total too small'
  54.         if self.total_length > piece_size * self.numpieces:
  55.             raise BTFailure, 'bad data in responsefile - total too big'
  56.         self.finished = finished
  57.         self.numactive = array('H', [0] * self.numpieces)
  58.         self.inactive_requests = [1] * self.numpieces
  59.         self.amount_inactive = self.total_length
  60.         self.endgame = False
  61.         self.have = Bitfield(self.numpieces)
  62.         self.waschecked = Bitfield(self.numpieces)
  63.         if self.numpieces < 32768:
  64.             typecode = 'h'
  65.         else:
  66.             typecode = 'l'
  67.         self.places = array(typecode, [NO_PLACE] * self.numpieces)
  68.         if not check_hashes:
  69.             self.rplaces = array(typecode, range(self.numpieces))
  70.             fastresume = True
  71.         else:
  72.             self.rplaces = self._load_fastresume(resumefile, typecode)
  73.             if self.rplaces is not None:
  74.                 fastresume = True
  75.             else:
  76.                 self.rplaces = array(typecode, [UNALLOCATED] * self.numpieces)
  77.                 fastresume = False
  78.         self.holepos = 0
  79.         self.stat_numfound = 0
  80.         self.stat_numflunked = 0
  81.         self.stat_numdownloaded = 0
  82.         self.stat_active = {}
  83.         self.stat_new = {}
  84.         self.stat_dirty = {}
  85.         self.download_history = {}
  86.         self.failed_pieces = {}
  87.  
  88.         if self.numpieces == 0:
  89.             return
  90.         targets = {}
  91.         total = 0
  92.         if not fastresume:
  93.             for i in xrange(self.numpieces):
  94.                 if self._waspre(i):
  95.                     self.rplaces[i] = ALLOCATED
  96.                     total += 1
  97.                 else:
  98.                     targets[hashes[i]] = i
  99.         if total and check_hashes:
  100.             statusfunc('checking existing file', 0)
  101.         def markgot(piece, pos):
  102.             if self.have[piece]:
  103.                 if piece != pos:
  104.                     return
  105.                 self.rplaces[self.places[pos]] = ALLOCATED
  106.                 self.places[pos] = self.rplaces[pos] = pos
  107.                 return
  108.             self.places[piece] = pos
  109.             self.rplaces[pos] = piece
  110.             self.have[piece] = True
  111.             self.amount_left -= self._piecelen(piece)
  112.             self.amount_inactive -= self._piecelen(piece)
  113.             self.inactive_requests[piece] = None
  114.             if not fastresume:
  115.                 self.waschecked[piece] = True
  116.             self.stat_numfound += 1
  117.         lastlen = self._piecelen(self.numpieces - 1)
  118.         partials = {}
  119.         for i in xrange(self.numpieces):
  120.             if not self._waspre(i):
  121.                 if self.rplaces[i] != UNALLOCATED:
  122.                     raise BTFailure("--check_hashes 0 or fastresume info "
  123.                                     "doesn't match file state (missing data)")
  124.                 continue
  125.             elif fastresume:
  126.                 t = self.rplaces[i]
  127.                 if t >= 0:
  128.                     markgot(t, i)
  129.                     continue
  130.                 if t == UNALLOCATED:
  131.             #Changed to accomodate pre-allocation
  132.             continue
  133.                 if t == ALLOCATED:
  134.                     continue
  135.                 if t!= FASTRESUME_PARTIAL:
  136.                     raise BTFailure("Bad fastresume info (illegal value)")
  137.                 data = self.storage.read(self.piece_size * i,
  138.                                          self._piecelen(i))
  139.                 self._check_partial(i, partials, data)
  140.                 self.rplaces[i] = ALLOCATED
  141.             else:
  142.                 data = self.storage.read(piece_size * i, self._piecelen(i))
  143.                 sh = sha(buffer(data, 0, lastlen))
  144.                 sp = sh.digest()
  145.                 sh.update(buffer(data, lastlen))
  146.                 s = sh.digest()
  147.                 if s == hashes[i]:
  148.                     markgot(i, i)
  149.                 elif s in targets and self._piecelen(i) == self._piecelen(targets[s]):
  150.                     markgot(targets[s], i)
  151.                 elif not self.have[self.numpieces - 1] and sp == hashes[-1] and (i == self.numpieces - 1 or not self._waspre(self.numpieces - 1)):
  152.                     markgot(self.numpieces - 1, i)
  153.                 else:
  154.                     self._check_partial(i, partials, data)
  155.                 statusfunc(fractionDone = 1 - self.amount_left /
  156.                            self.total_length)
  157.             if flag.isSet():
  158.                 return
  159.         self.amount_left_with_partials = self.amount_left
  160.         for piece in partials:
  161.             if self.places[piece] < 0:
  162.                 pos = partials[piece][0]
  163.                 self.places[piece] = pos
  164.                 self.rplaces[pos] = piece
  165.                 self._make_partial(piece, partials[piece][1])
  166.         for i in xrange(self.numpieces):
  167.             if self.rplaces[i] != UNALLOCATED:
  168.                 self.storage.allocated(piece_size * i, self._piecelen(i))
  169.             if self.have[i]:
  170.                 self.storage.downloaded(piece_size * i, self._piecelen(i))
  171.  
  172.     def _waspre(self, piece):
  173.         return self.storage.was_preallocated(piece * self.piece_size, self._piecelen(piece))
  174.  
  175.     def _piecelen(self, piece):
  176.         if piece < self.numpieces - 1:
  177.             return self.piece_size
  178.         else:
  179.             return self.total_length - piece * self.piece_size
  180.  
  181.     def _check_partial(self, pos, partials, data):
  182.         index = None
  183.         missing = False
  184.         marklen = len(self.partial_mark)+4
  185.         for i in xrange(0, len(data) - marklen,
  186.                         self.config['download_slice_size']):
  187.             if data[i:i+marklen-4] == self.partial_mark:
  188.                 ind = toint(data[i+marklen-4:i+marklen])
  189.                 if index is None:
  190.                     index = ind
  191.                     parts = []
  192.                 if ind >= self.numpieces or ind != index:
  193.                     return
  194.                 parts.append(i)
  195.             else:
  196.                 missing = True
  197.         if index is not None and missing:
  198.             i += self.config['download_slice_size']
  199.             if i < len(data):
  200.                 parts.append(i)
  201.             partials[index] = (pos, parts)
  202.  
  203.     def _make_partial(self, index, parts):
  204.         length = self._piecelen(index)
  205.         l = []
  206.         self.inactive_requests[index] = l
  207.         x = 0
  208.         self.amount_left_with_partials -= self._piecelen(index)
  209.         self.download_history[index] = {}
  210.         request_size = self.config['download_slice_size']
  211.         for x in xrange(0, self._piecelen(index), request_size):
  212.             partlen = min(request_size, length - x)
  213.             if x in parts:
  214.                 l.append((x, partlen))
  215.                 self.amount_left_with_partials += partlen
  216.             else:
  217.                 self.amount_inactive -= partlen
  218.                 self.download_history[index][x] = None
  219.         self.stat_dirty[index] = 1
  220.  
  221.     def _initalloc(self, pos, piece):
  222.         assert self.rplaces[pos] < 0
  223.         assert self.places[piece] == NO_PLACE
  224.         p = self.piece_size * pos
  225.         length = self._piecelen(pos)
  226.         if self.rplaces[pos] == UNALLOCATED:
  227.             self.storage.allocated(p, length)
  228.         self.places[piece] = pos
  229.         self.rplaces[pos] = piece
  230.         # "if self.rplaces[pos] != ALLOCATED:" to skip extra mark writes
  231.         mark = self.partial_mark + tobinary(piece)
  232.         mark += chr(0xff) * (self.config['download_slice_size'] - len(mark))
  233.         mark *= (length - 1) // len(mark) + 1
  234.         self.storage.write(p, buffer(mark, 0, length))
  235.  
  236.     def _move_piece(self, oldpos, newpos):
  237.         assert self.rplaces[newpos] < 0
  238.         assert self.rplaces[oldpos] >= 0
  239.         data = self.storage.read(self.piece_size * oldpos,
  240.                                  self._piecelen(newpos))
  241.         self.storage.write(self.piece_size * newpos, data)
  242.         if self.rplaces[newpos] == UNALLOCATED:
  243.             self.storage.allocated(self.piece_size * newpos, len(data))
  244.         piece = self.rplaces[oldpos]
  245.         self.places[piece] = newpos
  246.         self.rplaces[oldpos] = ALLOCATED
  247.         self.rplaces[newpos] = piece
  248.         if not self.have[piece]:
  249.             return
  250.         data = data[:self._piecelen(piece)]
  251.         if sha(data).digest() != self.hashes[piece]:
  252.             raise BTFailure('data corrupted on disk - '
  253.                             'maybe you have two copies running?')
  254.  
  255.     def _get_free_place(self):
  256.         while self.rplaces[self.holepos] >= 0:
  257.             self.holepos += 1
  258.         return self.holepos
  259.  
  260.     def get_amount_left(self):
  261.         return self.amount_left
  262.  
  263.     def do_I_have_anything(self):
  264.         return self.amount_left < self.total_length
  265.  
  266.     def _make_inactive(self, index):
  267.         length = self._piecelen(index)
  268.         l = []
  269.         x = 0
  270.         request_size = self.config['download_slice_size']
  271.         while x + request_size < length:
  272.             l.append((x, request_size))
  273.             x += request_size
  274.         l.append((x, length - x))
  275.         self.inactive_requests[index] = l
  276.  
  277.     def _load_fastresume(self, resumefile, typecode):
  278.         if resumefile is not None:
  279.             try:
  280.                 r = array(typecode)
  281.                 r.fromfile(resumefile, self.numpieces)
  282.                 return r
  283.             except Exception, e:
  284.                 self.errorfunc(WARNING, "Couldn't read fastresume data: " +
  285.                                str(e))
  286.         return None
  287.  
  288.     def write_fastresume(self, resumefile):
  289.         for i in xrange(self.numpieces):
  290.             if self.rplaces[i] >= 0 and not self.have[self.rplaces[i]]:
  291.                 self.rplaces[i] = FASTRESUME_PARTIAL
  292.         self.rplaces.tofile(resumefile)
  293.  
  294.     def get_have_list(self):
  295.         return self.have.tostring()
  296.  
  297.     def do_I_have(self, index):
  298.         return self.have[index]
  299.  
  300.     def do_I_have_requests(self, index):
  301.         return not not self.inactive_requests[index]
  302.  
  303.     def new_request(self, index):
  304.         # returns (begin, length)
  305.         if self.inactive_requests[index] == 1:
  306.             self._make_inactive(index)
  307.         self.numactive[index] += 1
  308.         self.stat_active[index] = 1
  309.         if index not in self.stat_dirty:
  310.             self.stat_new[index] = 1
  311.         rs = self.inactive_requests[index]
  312.         r = min(rs)
  313.         rs.remove(r)
  314.         self.amount_inactive -= r[1]
  315.         if self.amount_inactive == 0:
  316.             self.endgame = True
  317.         return r
  318.  
  319.     def piece_came_in(self, index, begin, piece, source = None):
  320.         if self.places[index] < 0:
  321.             if self.rplaces[index] == ALLOCATED:
  322.                 self._initalloc(index, index)
  323.             else:
  324.                 n = self._get_free_place()
  325.                 if self.places[n] >= 0:
  326.                     oldpos = self.places[n]
  327.                     self._move_piece(oldpos, n)
  328.                     n = oldpos
  329.                 if self.rplaces[index] < 0 or index == n:
  330.                     self._initalloc(n, index)
  331.                 else:
  332.                     self._move_piece(index, n)
  333.                     self._initalloc(index, index)
  334.  
  335.         if index in self.failed_pieces:
  336.             old = self.storage.read(self.places[index] * self.piece_size +
  337.                                     begin, len(piece))
  338.             if old != piece:
  339.                 self.failed_pieces[index][self.download_history[index][begin]]\
  340.                     = None
  341.         self.download_history.setdefault(index, {})
  342.         self.download_history[index][begin] = source
  343.  
  344.         self.storage.write(self.places[index] * self.piece_size + begin, piece)
  345.         self.stat_dirty[index] = 1
  346.         self.numactive[index] -= 1
  347.         if self.numactive[index] == 0:
  348.             del self.stat_active[index]
  349.         if index in self.stat_new:
  350.             del self.stat_new[index]
  351.         if not self.inactive_requests[index] and not self.numactive[index]:
  352.             del self.stat_dirty[index]
  353.             if sha(self.storage.read(self.piece_size * self.places[index], self._piecelen(index))).digest() == self.hashes[index]:
  354.                 self.have[index] = True
  355.                 self.storage.downloaded(index * self.piece_size,
  356.                                         self._piecelen(index))
  357.                 self.inactive_requests[index] = None
  358.                 self.waschecked[index] = True
  359.                 self.amount_left -= self._piecelen(index)
  360.                 self.stat_numdownloaded += 1
  361.                 for d in self.download_history[index].itervalues():
  362.                     if d is not None:
  363.                         d.good(index)
  364.                 del self.download_history[index]
  365.                 if index in self.failed_pieces:
  366.                     for d in self.failed_pieces[index]:
  367.                         if d is not None:
  368.                             d.bad(index)
  369.                     del self.failed_pieces[index]
  370.                 if self.amount_left == 0:
  371.                     self.finished()
  372.             else:
  373.                 self.data_flunked(self._piecelen(index), index)
  374.                 self.inactive_requests[index] = 1
  375.                 self.amount_inactive += self._piecelen(index)
  376.                 self.stat_numflunked += 1
  377.  
  378.                 self.failed_pieces[index] = {}
  379.                 allsenders = {}
  380.                 for d in self.download_history[index].itervalues():
  381.                     allsenders[d] = None
  382.                 if len(allsenders) == 1:
  383.                     culprit = allsenders.keys()[0]
  384.                     if culprit is not None:
  385.                         culprit.bad(index, bump = True)
  386.                     del self.failed_pieces[index] # found the culprit already
  387.                 return False
  388.         return True
  389.  
  390.     def request_lost(self, index, begin, length):
  391.         self.inactive_requests[index].append((begin, length))
  392.         self.amount_inactive += length
  393.         self.numactive[index] -= 1
  394.         if not self.numactive[index] and index in self.stat_active:
  395.             del self.stat_active[index]
  396.             if index in self.stat_new:
  397.                 del self.stat_new[index]
  398.  
  399.     def get_piece(self, index, begin, length):
  400.         if not self.have[index]:
  401.             return None
  402.         if not self.waschecked[index]:
  403.             if sha(self.storage.read(self.piece_size * self.places[index], self._piecelen(index))).digest() != self.hashes[index]:
  404.                 raise BTFailure, 'told file complete on start-up, but piece failed hash check'
  405.             self.waschecked[index] = True
  406.         if begin + length > self._piecelen(index):
  407.             return None
  408.         return self.storage.read(self.piece_size * self.places[index] + begin, length)
  409.